'''
Description: 
Author: Egrt
Date: 2022-10-24 17:07:38
LastEditors: Egrt
LastEditTime: 2022-10-24 18:50:36
'''
#------------------------------------------------#
#   进行训练前需要利用这个文件生成人脸的深度信息
#------------------------------------------------#
import os
import torch
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
from nets.prnet import PRNet
from PIL import Image
from utils.utils import cvtColor, preprocess_input, resize_image

#------------------------------------------------------#
#   载入PRNet模型及其权重
#------------------------------------------------------#
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
prnet = PRNet(3, 3)
pretrained_dict = torch.load('model_data/prnet_pytorch.pth', map_location = device)
prnet.load_state_dict(pretrained_dict)
prnet = prnet.eval()
prnet = prnet.cuda()

if __name__ == "__main__":
    #---------------------#
    #   训练集所在的路径
    #---------------------#
    datasets_path   = "datasets"

    types_name      = os.listdir(datasets_path)
    types_name      = sorted(types_name)

    for cls_id, type_name in enumerate(tqdm(types_name)):
        photos_path = os.path.join(datasets_path, type_name)
        if not os.path.isdir(photos_path):
            continue
        photos_name = os.listdir(photos_path)

        for photo_name in photos_name:
            if photo_name.endswith('.jpg'):
                image_path = os.path.join(os.path.abspath(datasets_path), type_name, photo_name)
                image = cvtColor(Image.open(image_path))
                image_256 = resize_image(image, [256, 256], letterbox_image = True)
                image_256 = np.transpose(preprocess_input(np.array(image_256, dtype='float32')), (2, 0, 1))
                image_256 = image_256[np.newaxis,:]
                image_256 = torch.from_numpy(np.array(image_256)).type(torch.FloatTensor)
                image_256 = image_256.cuda()
                image_features_3d = prnet(image_256)
                image_features_3d = F.interpolate(image_features_3d, size=(112, 112), mode='bicubic', align_corners=True)
                image_features_3d = image_features_3d.squeeze().cpu().detach().numpy()
                image_3d_path = os.path.join(os.path.abspath(datasets_path), type_name, photo_name).replace('jpg', 'npy')
                np.save(image_3d_path, image_features_3d)


            